import numpy as np
import pandas as pd
import warnings
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import norm

warnings.filterwarnings('ignore')

__all__ = ['FTProcessor', 'Regression']


class FTProcessor:

    @staticmethod
    def standardize_cap(x: np.ndarray) -> np.ndarray:
        x = (x - x.mean()) / x.std()
        mu = np.mean(x)
        std = np.std(x)
        ceil = mu + 3 * std
        floor = mu - 3 * std
        x[x > ceil] = ceil
        x[x < floor] = floor
        return x

    @staticmethod
    def sort_cap(x: np.ndarray) -> np.ndarray:  # see the chart under same directory
        x = np.argsort(np.argsort(x)) + 1  # descending order: range[1, n]
        x = x / (x.max() + 1)  # x / (n+1) , sum = 0.5 * n
        x = norm.ppf(x)  # norm.ppf() (inverse of cdf)
        x[x > 3] = 3
        x[x < -3] = -3
        return x

    @staticmethod  # 截面填充:一般适用于计算好某日因子值后在截面填充中位数,数据结构往往是['date', 'code', feature],date.unique==1.
    def fillna_cs(df: pd.DataFrame) -> pd.DataFrame:
        # assert len(df['date'].unique()) != 1, f"INFO | There are more than one date in this df."
        if df.set_index(['date', 'code']).isna().sum().sum() != 0:
            median = df.set_index(['date', 'code']).median()[0]
            df = df.fillna(median)
        return df

    @staticmethod  # 时序填充:一般适用于时序提取feature矩阵时,对缺失值使用上期值填充
    def fillna_ts_expand(df: pd.DataFrame) -> pd.DataFrame:
        df_copy = df.copy(deep=False)
        col_name = df_copy.columns[-1]
        # transform to matrix
        df = df.set_index(['date', 'code']).unstack().sort_index(ascending=False).droplevel(0, axis=1)  # 时间降序排列
        if df.isna().sum().sum() != 0:
            df.fillna(method='ffill', inplace=True)
            df = df.stack().to_frame(name=col_name).reset_index()
            df_copy = df_copy.merge(df, on=['date', 'code'], how='left', suffixes=('_pre', ''))
            df_copy.drop(columns=[col_name+'_pre'], inplace=True)
            return df_copy
        else:
            return df_copy


    @staticmethod
    def expand(df: pd.DataFrame) -> pd.DataFrame:
        # transform to matrix
        df = df.set_index(['date', 'code']).unstack().sort_index(ascending=False).droplevel(0, axis=1)  # 时间降序排列
        return df

    @staticmethod
    def check_codes_loss(df: pd.DataFrame, codes: list):
        df_codes = df.columns.tolist()
        if df_codes == codes:
            return df
        else:
            return pd.DataFrame(index=codes).join(df.T, how='left').T

    # outliers
    @staticmethod
    def simple_mad(arr: np.array) -> np.array:
        arr[arr > 3] = 3
        arr[arr < -3] = -3
        return arr

    @staticmethod  # one-step Gram-Schmidt Procedure
    def orthogonal(x: np.ndarray, y: np.ndarray, sort_cap: bool = False) -> np.ndarray:
        """
        params: y -> target, np.ndarray [n, ]
        params: x -> np.ndarray [n, ]
        return: res -> np.ndarray [n, ]
        """
        x = np.matrix(x).T
        y = np.matrix(y).T
        b = ((x.T - x.mean()) * y) / (x.T * x - x.mean() * x.mean() * len(x))
        b0 = y.mean() - b * x.mean()
        res = np.array(y - b0 - x * b).squeeze()
        """
        params: y -> target, n*1 matrix
        params: x -> n*1 matrix
        return: res -> n*1 matrix
        """
        if sort_cap:
            return FTProcessor.sort_cap(res)
        else:
            return res

    @staticmethod
    def ewm(T: int, halflife: int) -> np.ndarray:
        t = np.arange(T) + 1
        c = (0.5 ** (1 / halflife)) ** (T - t)
        c = c / c.sum()
        return c

class Regression:

    @staticmethod
    def wls(x, y, w, add_constant=False):
        # check shape
        if len(x.shape) == 1:
            x = x.reshape(-1, 1)
        # add constants to x
        if add_constant:
            ones = np.ones([x.shape[0], 1])
            x = np.hstack([ones, x])

        # b_hat estimation
        w_inv = np.linalg.inv(np.diag(w))  # highlight: barra将流动市值的平方根作为残差的异方差。
        x_t_w_inv_x = np.matmul(np.matmul(np.transpose(x), w_inv), x)
        try:
            x_t_w_inv_x_inv = np.linalg.inv(x_t_w_inv_x)
        except np.linalg.LinAlgError:
            x_t_w_inv_x_inv = np.linalg.pinv(x_t_w_inv_x)
        x_t_w_inv_y = np.matmul(np.matmul(np.transpose(x), w_inv), y)
        b_hat = np.matmul(x_t_w_inv_x_inv, x_t_w_inv_y)
        epsilon_hat = y - np.matmul(x, b_hat)
        return [b_hat, epsilon_hat]

    @staticmethod
    def ols(x, y, add_constant=False):
        w = np.ones(x.shape[0])
        return Regression.wls(x, y, w, add_constant=add_constant)